# -*- coding: utf-8 -*-
# $Id: udineGenUtils2.py 14 2010-02-05 04:05:12Z leo.monash.uni $
# Copyright 2010, the Regents of the University of Arizona
# You may use this file under the terms of the LGPL.
# Written by Leo Lopes 
# from the original generator written by Jakub Mareček
#
"""You can get the code, as well as information on which papers to cite
if you use this code, at http://code.google.com/p/udinettgen

To run, make sure that combinedITC07graphs.pickle is in the same
directory, and run 'python udinettgen2.py 100 70' where 
100 is the number of events (lectures acrosss courses) and
70 is the occupancy (rooms*periods - lectures)
There is also an example of how to use the classes in your own code
at the bottom of this file, along with a comprehensive unit test if
you have a bash shell available.
"""

import math, string, glob
import pickle, sys
from math import ceil
from random import seed, random, choice, randrange, getstate, setstate, randint, uniform
import networkx

class TimetablingInstance:
  "A class storing properties of a (randomly generated) instance of a course timetabling problem"
  # The connected components from all the curricula conflict graphs 
  # and other metatada from the 21 itc2007 instances. 
  itc07Graph, itc07enrollment = pickle.load(open('combinedITC07graphs.pickle'))
  itc07components = networkx.connected_component_subgraphs(itc07Graph)
  oneRoomEventDistr = ((0,14),(1,3),(2,1),(4,1),(7,1),(9,1))
  
  def __init__(self, events, occupancy,nameConvention="Name: Random_%iEvents_Occupancy%i\n"):
    #seed(1) # uncomment for debugging.
    self.pEvents = events
    self.pOccupancy = float(occupancy)/100
    self.nameConvention = nameConvention

    self.pDays = 5
    self.pPeriodsPerDay = 5
    self.pPeriods = self.pDays * self.pPeriodsPerDay 

    self.pMinEventsPerCourse = 2
    self.pMaxEventsPerCourse = 5
    self.pCourses = 0            # reset in genCourses!
    self.pMinDaysFactor = 1.125
    self.enrollment = {}

    self.pTeachers = int(self.pEvents / 6)
    self.pMinCoursesPerTeacher = 1
    self.pMaxCoursesPerTeacher = 3

    #self.pCurricula = int(self.pEvents / 10)
    self.pCurriculaToCoursesRatio = uniform(.4,1.2)
    self.pMinCoursesPerCurriculum = 2
    self.pMaxCoursesPerCurriculum = 6
    # curriculaGraph and associated data are for sampling.
    self.curriculaGraph = networkx.Graph() 
    self.curriculaGraphComponents = []

    self.pRooms = max(int(ceil( (self.pEvents / self.pPeriods) / self.pOccupancy)), 1)
    self.pDistinctRoomSizes = max(int(ceil(0.6 * self.pRooms)), 1)
    self.oneRoomEvents = choiceWithWeights(self.oneRoomEventDistr)

    self.pSmallestClassroom = 40
    self.pLargestClassroom = 350

    self.pConstraints = int(self.pEvents * 1.2)

    self.intRoomSizes = []
    self.intCourseIds = []
    

  def genRooms(self):
    outRooms = "\nROOMS:\n"
    distinct = []
    while len(distinct) < self.pDistinctRoomSizes:
      newSize = randrange(self.pSmallestClassroom, self.pLargestClassroom, 5)
      if newSize not in distinct: distinct.append(newSize)
    distinct.sort()
    complete = distinct[:]
    # make sure it is possible to generate one room events if required.
    if self.oneRoomEvents and len(distinct) > 1: distinct.pop()
    for j in range(self.pRooms - self.pDistinctRoomSizes):
      complete.append(choice(distinct))
    complete.sort()
    self.intRoomSizes = complete[:] # used later to match with class sizes.
    for room, capacity in zip(range(self.pRooms), complete):
      outRooms += "R%02i\t%i\n" % (room + 1, capacity)
    return outRooms
    
  def genCourses(self):
    outCourses = "\nCOURSES:\n"
    eventsCounts = []
    minDaysSpecs = []
    taughtEvents = 0
    while taughtEvents < self.pEvents:
      ev = min(randrange(self.pMinEventsPerCourse, self.pMaxEventsPerCourse), self.pEvents - taughtEvents)
      eventsCounts.append(ev)
      minDaysSpecs.append(int(ev/self.pMinDaysFactor))
      taughtEvents += ev
    self.pCourses = len(eventsCounts)
    self.pCurricula = int(self.pCurriculaToCoursesRatio*self.pCourses)
    taught = self.pMinCoursesPerTeacher * range(self.pTeachers)
    teachers = set(range(self.pTeachers))
    while len(taught) < self.pCourses:
      teacher = teachers.pop()
      for i in range(max(self.pMaxCoursesPerTeacher - self.pMinCoursesPerTeacher, len(taught) - self.pCourses)):
        taught.append(teacher)
    # How course sizes are generated determines the number of 1-room events. Handle 1-room events seperately.
    students = []
    minStu,maxStu = int(self.pSmallestClassroom/2),self.intRoomSizes[-1]
    if len(self.intRoomSizes) > 1: maxStu = self.intRoomSizes[-2]
    coursesToGenerate = self.pCourses
    # It is only possible to explicitly control the number of one room events 
    # if there is more than one room available
    if self.oneRoomEvents and len(self.intRoomSizes)>1: 
#      sys.stderr.write("%d %s\n" % (self.oneRoomEvents,self.intRoomSizes))
      oneRoomEventsGenerated = min(self.oneRoomEvents,self.pCourses)
      students.extend([randrange(self.intRoomSizes[-2],self.intRoomSizes[-1]) for i in range(oneRoomEventsGenerated)])
      maxStu = self.intRoomSizes[-2]
      coursesToGenerate -= oneRoomEventsGenerated
    students.extend([randrange(minStu, maxStu) for j in range(coursesToGenerate)])
    for course, teacher, eventsCount, minDays, capacity in zip(range(self.pCourses), taught, eventsCounts, minDaysSpecs, students):
      courseid = "Course%03i" % (course + 1)
      self.intCourseIds.append(courseid)
      # keep track of enrollment to use for sampling curricula later.
      self.enrollment[courseid] = capacity
      outCourses += "%s t%02i %i %i %i\n" % (courseid, teacher, eventsCount, minDays, capacity)
    return outCourses

  def genUnavailability(self):
    outUnavail = "\nUNAVAILABILITY_CONSTRAINTS:\n"
    unavail = []
    while len(unavail) < self.pConstraints:
      course = choice(self.intCourseIds)
      day = randrange(0, self.pDays)
      period = randrange(0, self.pPeriodsPerDay)
      if random() > 0.5:
        for p in range(period, self.pPeriodsPerDay): unavail.append((course, day, p))
      else:
        for p in range(0, period+1): unavail.append((course, day, p))
    unavail = sorted(list(set(unavail))) # May result in the number of constraints being cut, so...
    self.pConstraints = len(unavail) # ... update the number of constraints.
    for course, day, period in unavail:
      outUnavail += "%s %i %i\n" % (course, day, period)
    return outUnavail

  def genCurricula(self):

    def sampleGraph():
      """Produce a graph based on the itc2007 connected components 
      where the number of nodes matches the number of courses selected above."""
      # Only interested in components longer than shortest curriculum.
      largerThanShortestCurriculum = [c for c in self.itc07components if len(c.nodes()) >= self.pMinCoursesPerCurriculum]
      components, nCourses, = [], 0
      while(nCourses < self.pCourses):
        lastComponent = choice(largerThanShortestCurriculum)
        components.append(lastComponent)
        nCourses = sum([len(c) for c in components])
      # find the component which has to be modified the least to fit.
      # Try to keep the large components as intact as possible
      # since they are more likely to constrain the problem.
      if nCourses > self.pCourses:
        excess = nCourses - self.pCourses
        cutSize, toCut = sorted([(len(c)-excess,c) for c in components if len(c) >= excess ])[0]
        del components[components.index(toCut)]
        # if we've cut too many classes by removing toCut, add back a connected subgraph
        if cutSize:
          components.append(toCut.subgraph(self.curriculumFromWeighedWalk(cutSize,toCut,self.itc07enrollment)))
      # rename the nodes with enrollments in the same order. 
      fullGraph = self.itc07Graph.subgraph(sum([c.nodes() for c in components],[]))
      itcNodesBySize = sorted(fullGraph.nodes(),key=lambda x: self.itc07enrollment[x])
      thisInstanceNodesBySize = sorted(self.enrollment,key = lambda x: self.enrollment[x])
      translation = dict(zip(itcNodesBySize,thisInstanceNodesBySize))
      self.curriculaGraph = networkx.Graph()
      for (old,new) in translation.iteritems():
        self.curriculaGraph.add_node(new)
        for v in fullGraph.neighbors(old): self.curriculaGraph.add_edge(new,translation[v])
      # compute this helpful quantity only once
      self.curriculaGraphComponents = networkx.connected_component_subgraphs(self.curriculaGraph)
      # now make sure that no impossible curricula are generated
      self.pMaxCoursesPerCurriculum = min(self.pMaxCoursesPerCurriculum,max(len(G) for G in self.curriculaGraphComponents))
      
    outCurr = "\nCURRICULA:\n"
    curr = []
    # Initialize sample graph. Graph must have at least one component large enough 
    # to generate the largest curriculum.
    sampleGraph()
    # first make sure every course is used at least once
    unusedCourses = set(self.curriculaGraph.nodes())
    nCurricula = 0
    while unusedCourses:
      start = choiceWithWeights([(c,self.enrollment[c]) for c in unusedCourses])
      component = [g for g in self.curriculaGraphComponents if start in g.nodes()][0]
      nCourses = randrange(min(self.pMinCoursesPerCurriculum,len(component)),min(self.pMaxCoursesPerCurriculum,len(component))+1)
      courses = self.curriculumFromWeighedWalk(nCourses,component,self.enrollment,start)
      curr.append((nCurricula,courses))
      unusedCourses.difference_update(courses)
      nCurricula += 1
    # now if there are any curricula left add them.
    for i in range(nCurricula,self.pCurricula):
      nCourses = randrange(self.pMinCoursesPerCurriculum,self.pMaxCoursesPerCurriculum+1)
      #sys.stderr.write("%d %d %d\n" % (nCourses,self.pMinCoursesPerCurriculum,self.pMaxCoursesPerCurriculum+1))
      bigEnoughComponents = [g for g in self.curriculaGraphComponents if len(g.nodes()) >= nCourses]
      component = choice(bigEnoughComponents)
      courses = self.curriculumFromWeighedWalk(nCourses,component,self.enrollment)
      curr.append((i, courses))
    self.pCurricula = len(curr)  
    for currid, courses in curr:
      courses = sorted(list(courses))
      outCurr += "Curr%003i %i %s\n" % (currid+1, len(courses), string.join(courses, " "))
    return outCurr

  def genHeader(self,name=None):
    if name: outHeader = "Name: %s\n" % name
    else: outHeader = self.nameConvention % (self.pEvents, int(self.pOccupancy * 100))
    outHeader += "Courses: %i\n" % self.pCourses
    outHeader += "Rooms: %i\n" % self.pRooms
    outHeader += "Days: %i\n" % self.pDays
    outHeader += "Periods_per_day: %i\n" % self.pPeriodsPerDay
    outHeader += "Curricula: %i\n" % self.pCurricula
    outHeader += "Constraints: %i\n" % self.pConstraints
    return outHeader

  def genInstanceDef(self,name=None):
    rooms = self.genRooms()
    courses = self.genCourses()
    curr = self.genCurricula()
    unavail = self.genUnavailability()
    header = self.genHeader(name)
    return header + courses + rooms + curr + unavail + "\nEND."

  def curriculumFromWeighedWalk(self,length, graph=None,weights=None,initialNode=None):
    """Returns a random set of courses of length 'length' connected through a graph.
    The list is obtained from a random walk where the probability of a step is
    obtained by the relative weights of nodes. Thus this procedure can cycle for ever.
    'graph' is a networkx.Graph. 
    'weights' is a dict with entries for each node in 'graph'
    """
    # random walks can take a long time to complete if 
    # some courses have very low enrollment and length is close to the size of the component.
    # The following lines help in those cases.
    if initialNode and initialNode not in graph.nodes(): 
      raise KeyError, "No such node %s. Nodes: %s" % (initialNode,graph.nodes())
    if length == len(graph): seq = set(graph.nodes())
    else:
      if initialNode: last = initialNode
      else: last = choiceWithWeights([(v,weights[v]) for v in graph.nodes()])
      seq = set([last])
      while len(seq) < length:
        last = choiceWithWeights([(v,weights[v]) for v in graph.neighbors(last)])
        seq.add(last)
    return seq

def choiceWithWeights(elemFreqTuples):
  """Like choice but takes into account weights."""
  X = random()*sum([f for (i,f) in elemFreqTuples])
  s = 0.
  for (i,f) in elemFreqTuples:
    s += f
    if X < s: break
  return i


if __name__ == '__main__':
  try:
    if len(sys.argv) == 3:
      state = getstate()
      print TimetablingInstance(int(sys.argv[1]),int(sys.argv[2])).genInstanceDef()
    elif len(sys.argv) == 2:
      state,argv,e = pickle.load(open(sys.argv[1]))
      print "Handling exception: \n\n%s\n\n" % e
      setstate(state)
      print TimetablingInstance(int(argv[1]),int(argv[2])).genInstanceDef()
  except:
    pickle.dump([state,sys.argv,sys.exc_info()[1]],open('%s-%9d.core' % ("-".join(sys.argv[1:]),randint(0,999999999)),'w'))
    raise

unitTest = """
#bash script:
for events in 20 40 50 70 100 125 150 200 400; do 
  for occupancy in 30 50 70 80 90; do 
    for i in `seq 1 100`; do 
      echo -n ${events} ${occupancy} ${i}", "; 
      python udineGenUtils2.py ${events} ${occupancy} > ${events}Events${occupancy}PctOccupancy-${i}.ctt; 
    done 
  done 
done
"""
